package tool;

import org.apache.commons.io.IOUtils;
import org.tensorflow.Graph;
import org.tensorflow.Session;
import org.tensorflow.Tensor;

import java.awt.image.BufferedImage;
import java.io.FileInputStream;
import java.nio.FloatBuffer;

/**验证码识别类
 * Create by zekdot on 19-12-25.
 */
public class CodeRecog {
    private static char RES_CHAR[]; //可能的结果

    static{
        RES_CHAR=new char[36];//有36种可能结果
        for(int i=0;i<10;i++){
            RES_CHAR[i]= (char) (48+i);   //数字
        }
        for(int i=10;i<36;i++){
            RES_CHAR[i]= (char) (i+87);   //字母
        }
    }
    /**根据输入获取输出
     *
     * @param modelPath 模型所在位置
     * @param input 图像数组
     * @return 图像对应的字符
     */
    private static char verifyOneCheckCode(float[] input,String modelPath){
        try{
            Graph graph = new Graph();  //新建一个计算图
            //导入图
            byte[] graphBytes = IOUtils.toByteArray(new
                    FileInputStream(modelPath + "model.pb"));
            graph.importGraphDef(graphBytes);
            FloatBuffer floatBuffer=FloatBuffer.allocate(240);  //申请一个240个double的空间
            long shape[]={1,240};   //新建维度数组
            floatBuffer.put(input,0,240);   //将数据放入数组
            floatBuffer.position(0);    //修改指针为开头
            Tensor<Float> tensor=Tensor.create(shape,floatBuffer);  //创建一个张量
            //根据图建立Session
            Session session = new Session(graph);   //新建一个会话
            Tensor<Float> result= (Tensor<Float>) session.runner().feed("x-input",tensor).fetch("predict").run().get(0);    //传入测试张量算出结果
            floatBuffer.position(0);    //设置当前指针为开头
            result.writeTo(floatBuffer);    //写入到缓存区
            floatBuffer.position(0);    //设置指针为开头
            float res[]=new float[36];  //申请一个结果数组
            floatBuffer.get(res,0,36);  //讲张量数据写入结果数组
            int index=0;//索引
            float pos=res[0];//可能性
            for(int i=1;i<res.length;i++){
                if(res[i]>pos){ //如果发现了更大的可能性
                    pos=res[i];
                    index=i;
                }
            }
            session.close();    //关闭会话
            tensor.close();;//关闭张量
            return RES_CHAR[index];//返回可能字符
        }catch (Exception e){
            e.printStackTrace();
        }
        return 0;
    }

    /** 根据图片对象识别内容
     *
     * @param modelPath 模型所在位置
     * @param image 图片
     * @return 图片对应的四个字符的字符串
     */
    public static String verifyCheckCode(BufferedImage image,String modelPath){
        StringBuffer buffer=new StringBuffer();
        float arrays[][]= PicTool.getArraysByPic(image);
        for(int i=0;i<4;i++){
            buffer.append(verifyOneCheckCode(arrays[i],modelPath));
        }
        return buffer.toString();
    }

}
